#!/usr/bin/env python3
# -*- coding:utf-8 -*-
"""
last edited at 2019-7-19
by Stone @ BlueNet

HSV空间颜色识别与圆形匹配识别乒乓球。
"""
from collections import deque
import cv2 as cv
import logging
from matplotlib import pyplot as plt
import numpy as np
import sys

from tools import performance_measure as pmea


DEBUG = True
DEV_CAMERA = r'..\\experiment\\images\\广工小车.mp4'  # device camera, which is to be used

# H, S for color recognition is enough
# HS_TH = ((15, 20), (1, 255))  # threshold of hue and saturation
# HS_LOWER = (10, 100)
# HS_UPPER = (20, 255)
HS_LOWER = (114, 100)
HS_UPPER = (206, 255)
SHAPE_TH = 0.15  # threshold to consider shape matching success
LEN_TH = 15  # threshold to filter the contours

IMG_SIZE = np.asarray((640, 480))  # the size of final image used to find the ball
# IMG_SIZE = np.asarray((400, 400))  # the size of final image used to find the ball
RESIZE_RATIO = 0.3  # the ratio to resize the input image
RADIUS_RANGE = (IMG_SIZE[0]/4, IMG_SIZE[0]/2)  # the range of circle's radius

PTH_SHAPE_CIRCLE = r'shape_circle.png'  # a shape used to match the circle, generated using OpenCV


def in_range_hsv(img_hsv, hsv_lower=HS_LOWER, hsv_upper=HS_UPPER):
    """
    基于HSV空间颜色阈值的图像二值化操作. 目前只用到H, S阈值, 而不使用V 的阈值,
    但是保留其位置以兼容未来的升级.

    :param img_hsv: the image in hsv color space to be threshold.
    :param hsv_lower: the lower threshold of hsv value. (currently using hs only)
    :param hsv_upper: the upper threshold of hsv value. (currently using hs only)
    :return: the bw image in np.uint8, the pixels in threshold is white, the pixels
        out of the threshold is black
    """
    img_wb = cv.bitwise_and(
        cv.inRange(img_hsv[:, :, 0], hsv_lower[0], hsv_upper[0]),
        cv.inRange(img_hsv[:, :, 1], hsv_lower[1], hsv_upper[1])
    )

    # debug
    # cv.imshow("bw", img_wb)

    return img_wb


def filter_shape(contours):
    """
    Filter the contours by length(quantity of points of the contour) and its shape(circle).
    :param contours: the contours to be filtered.
    :return: the filtered contours that may be a circle.
    """
    cnt_rslt = deque()  # result contours filtered

    # filter the contours by length and shape
    for cnt in contours:
        if len(cnt) > LEN_TH:
            if cv.matchShapes(cnt, STD_CNT_CIRCLE, cv.CONTOURS_MATCH_I1, 0) < SHAPE_TH:
                cnt_rslt.append(cnt)

        # debug
        # logging.info("len_of_cnt: %i" % len(cnt))
        # print 'shape matching: %.6f' % cv.matchShapes(cnt, STD_CNT_CIRCLE, cv.CONTOURS_MATCH_I1, 0)

    return cnt_rslt


def init():
    global STD_CNT_CIRCLE  # standard contour circle, in OpenCV contour format (x,y array)
    global counter  # time counter for measuring performance

    logging.basicConfig(level=logging.INFO)

    img_circle = cv.imread(PTH_SHAPE_CIRCLE, cv.IMREAD_GRAYSCALE)
    _, STD_CNT_CIRCLE, _ = cv.findContours(img_circle, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    STD_CNT_CIRCLE = STD_CNT_CIRCLE[0]

    counter = pmea.TimeCounter()

    logging.info("init done")


def main():
    src = 0  # 0: camera; others: file

    if src == 0:
        camera = cv.VideoCapture(DEV_CAMERA)
        if not camera.isOpened():
            logging.error("Failed to open camera.")
    else:
        frame = cv.imread(r'D:\AAA_BlueNet\BluePingPong\pingpong2.jpg')
        ret = True

    while True:

        if src == 0:
            ret, frame = camera.read()
        # img_raw = cv.resize(frame, tuple(IMG_SIZE))
        img_raw = frame

        if not ret:
            # fail to get image
            continue

        # color threshold
        img_hsv = cv.cvtColor(img_raw, cv.COLOR_BGR2HSV)
        img_bw = in_range_hsv(img_hsv)

        # shape filtering
        _, contours, _ = cv.findContours(img_bw, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
        cnt_targets = filter_shape(contours)

        for i in range(len(cnt_targets)):
            # draw contours of found balls
            cv.drawContours(img_raw, cnt_targets, i, (255, 0, 0), 3)

        if DEBUG:
            # draw all the contours
            [cv.drawContours(img_raw, contours, i, (0, 255, 0), 1) for i in range(len(contours))]
            logging.info("num of result: %i" % len(cnt_targets))  # debug
            logging.info("time used: %s" % str(counter.get_delta_count()))

            # img_bw_resized = cv.resize(img_bw, (50, 50))
            # # img_bw_resized = cv.resize(img_bw_resized, (800,800))
            # logging.info(str(
            #     len(cv.findContours(img_bw_resized, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)[1])
            # ))
            # # cv.imshow("bw-resize", img_bw_resized)

        cv.imshow("img_raw", img_raw)
        if cv.waitKey(100) == ord('q'):
            # do some cleaning
            camera.release()
            cv.destroyAllWindows()
            logging.info('Program exited.')

            sys.exit()


if __name__ == '__main__':
    init()
    main()
